# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
# ~~~~~~~ 使用LSB隐写术解码图像中的秘密消息 ~~~~~~~ #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

from skimage.io import imread, imsave
import sys
import wave

from DataHidingDemo.extra import bin_to_file, bin_to_mes

from PIL import PngImagePlugin
PngImagePlugin.MAX_TEXT_CHUNK = 500 * (1024**2)

def lsb_reveal(img, bits=1, out_name='output_file.txt'):
    # 找出图像的大小和频道数
    a, b = img.shape[0], img.shape[1]
    c = 3 if len(img.shape) == 3 else 1
    bitstream = []
    for i in range(a):
        for j in range(b): 
            for k in range(c): 
                # 获取像素的最后一位并将其添加到位流中 
                if c == 1: 
                    bitstream.extend(bin(img[i, j] % (2 ** bits))[2:].rjust(bits, "0")) 
                else: 
                    bitstream.extend(bin(img[i, j, k] % (2 ** bits))[2:].rjust(bits, "0"))
    return bin_to_file(bitstream, out_name)

def wav_lsb_extract(wav):
    cover_wav = wave.open(wav, mode='rb')
    frames = bytearray(cover_wav.readframes(cover_wav.getnframes()))

    len_bits = [int(f) % 2 for f in frames[:16]]
    len_string = ''.join(map(str, len_bits))
    mess_len = int(len_string, 2)
    # print(mess_len)
    message_bits = [int(f) % 2 for f in frames[16:mess_len + 16]]
    # print(message_bits)
    message_ex = []
    value = 0
    
    for i in range(len(message_bits)):
        if i % 8 == 0 and i != 0:
            message_ex.append(value)
            value = 0
        value |= message_bits[i] << i % 8
    message_ex.append(value)
    result = ''.join([chr(l) for l in message_ex])
    print(result)
    return result

def main():
    print("Enter the source path of the image to decode:")
    file = input()
    image = imread(file)
    lsb_reveal(image, 1)

if __name__ == '__main__':
	main()
